import streamlit as st
import torch
import cv2
import numpy as np
from PIL import Image

st.title("安全绳佩戴检测系统")

# 加载本地模型（假设使用 ultralytics YOLOv8）
from ultralytics import YOLO
model = YOLO('runs/detect/train9/weights/best.pt')  # 替换为你的本地模型路径
# model = YOLO('archive/results_yolov8n_100e/kaggle/working/runs/detect/train/weights/best.pt')  # 替换为你的本地模型路径



# 上传图片
uploaded_file = st.file_uploader("上传一张作业现场图片", type=["jpg", "png", "jpeg"])

def iou(boxA, boxB):
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    inter = max(0, xB - xA) * max(0, yB - yA)
    areaA = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
    areaB = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
    return inter / float(areaA + areaB - inter)

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    img_np = np.array(image)
    # ultralytics YOLOv8 推理
    results = model(img_np)
    print(results)
    boxes = results[0].boxes  # YOLOv8 结果
    cls = boxes.cls.cpu().numpy().astype(int)
    xyxy = boxes.xyxy.cpu().numpy()
    print(xyxy)
    print(cls)
    # 根据模型类别定义：names: [Hardhat, Mask, NO-Hardhat, NO-Mask, NO-Safety Vest, Person, Safety Cone, Safety Vest, machinery, vehicle]
    # 类别5是Person，类别7是Safety Vest（安全背心）
    persons = [xyxy[i] for i in range(len(cls)) if cls[i] == 5]  # Person
    hardhats = [xyxy[i] for i in range(len(cls)) if cls[i] == 0]  # Hardhat

    print('检测结果:')
    print(f'检测到 {len(persons)} 个人')
    print(f'检测到 {len(hardhats)} 个安全帽')
    
    # 为每个检测到的人绘制边界框和标签
    used_hardhats = set()
    for person_idx, person in enumerate(persons):
        x1, y1, x2, y2 = person
        # 取头部区域（上1/3）
        head_y2 = y1 + (y2 - y1) / 3
        head_box = [x1, y1, x2, head_y2]
        head_cx = (x1 + x2) / 2
        head_cy = (y1 + head_y2) / 2
        best_hat_idx = -1
        best_dist = float('inf')
        # 遍历所有未分配的安全帽
        for hat_idx, hat in enumerate(hardhats):
            if hat_idx in used_hardhats:
                continue
            iou_score = iou(head_box, hat)
            if iou_score > 0.1:
                hat_cx = (hat[0] + hat[2]) / 2
                hat_cy = (hat[1] + hat[3]) / 2
                dist = (head_cx - hat_cx) ** 2 + (head_cy - hat_cy) ** 2
                if dist < best_dist:
                    best_dist = dist
                    best_hat_idx = hat_idx
        if best_hat_idx != -1:
            has_hardhat = True
            used_hardhats.add(best_hat_idx)
        else:
            has_hardhat = False
        if has_hardhat:
            color = (0, 255, 0)  # 绿色 - 安全
            label = "Safe"
        else:
            color = (0, 0, 255)  # 红色 - 不安全
            label = "No Hardhat!"
        cv2.rectangle(img_np, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
        cv2.putText(img_np, label, (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)

    st.image(img_np, caption="检测结果", channels="BGR")
